#include "voxel_interpolation.h"

#include <stdio.h>
#include <math.h>
#include <vector>
#include <stack>
#include<iostream>

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

namespace py = pybind11;

int voxel_encode(float const* gt_lines_pts_ptr,
                 std::uint16_t const& gt_lines_num,
                 std::uint16_t const& gt_lines_length,
                 float const* voxels_size,
                 float const* pts_range,
                 std::uint16_t const& max_dense_pts_size,
                 std::uint16_t const& dense_pts_length,
                 std::uint16_t const& voxels_size_dim,
                 float* dense_pts_ptr){
  if(nullptr == voxels_size || nullptr == pts_range || nullptr == gt_lines_pts_ptr || nullptr == dense_pts_ptr){
     std::cout << "NULL pointer" << std::endl;
  }

  std::array<float, 3> voxels_size_;
  std::array<float, 6> pts_range_;
  std::array<float, 3> max_xyz_;

  for (std::uint16_t i{0U}; i < voxels_size_dim; ++i){
    if(std::fabs(voxels_size[i]) < EPS){
      std::cout << "voxels_size is equal to zero!" << std::endl;
    }

    voxels_size_[i] = voxels_size[i];
    pts_range_[i] = pts_range[i];
    pts_range_[voxels_size_dim+i] = pts_range[voxels_size_dim+i];
    max_xyz_[i] = std::round((pts_range[voxels_size_dim+i] - pts_range[i]) / voxels_size[i]);
  }

  std::uint16_t filled_size{0};
  std::uint16_t start_i{0};
  float id{0.0f};

  std::array<Point, 2048> ordered_pts;
  std::uint16_t ordered_pts_filled_size{0U};

  for (std::uint16_t i{0U}; i < gt_lines_num; ++i){
    float const& cur_id = gt_lines_pts_ptr[i * gt_lines_length + 0];
    float const& next_id = gt_lines_pts_ptr[(i+1) * gt_lines_length + 0];

    // line id changed or last point
    if(std::fabs(cur_id - next_id) > EPS || (i == gt_lines_num-1)){
      ordered_pts_filled_size = process_one_line(gt_lines_pts_ptr, start_i, i,
                                                  gt_lines_length, voxels_size_, pts_range_, ordered_pts);
      const float& cls = gt_lines_pts_ptr[start_i * gt_lines_length + 4];
      const float& type = gt_lines_pts_ptr[start_i * gt_lines_length + 3];

      float const& start_x = gt_lines_pts_ptr[start_i * gt_lines_length + 1];
      float const& start_y = gt_lines_pts_ptr[start_i * gt_lines_length + 2];
      float const& end_x = gt_lines_pts_ptr[i * gt_lines_length + 1];
      float const& end_y = gt_lines_pts_ptr[i * gt_lines_length + 2];
      if(((std::fabs(start_x - end_x) + std::fabs(start_y - end_y)) < EPS) &&
            ordered_pts_filled_size < ordered_pts.size() && ordered_pts_filled_size > 0){
        ordered_pts.at(ordered_pts_filled_size).x = ordered_pts.at(1).x;
        ordered_pts.at(ordered_pts_filled_size).y = ordered_pts.at(1).y;
        ordered_pts_filled_size++;
      }

      filled_size = fill_output(ordered_pts, ordered_pts_filled_size, id, cls, type, filled_size,
                                  max_xyz_, max_dense_pts_size, dense_pts_length, dense_pts_ptr);
      start_i = i + 1;
      ordered_pts_filled_size = 0U;
    }
  }
  return filled_size;
}

std::uint16_t fill_output(std::array<Point, 2048> const& ordered_pts,
                          std::uint16_t const& ordered_pts_filled_size,
                          float& id,
                          float const& cls,
                          float const& type,
                          std::uint16_t& filled_size,
                          std::array<float, 3> const& max_xyz_,
                          std::uint16_t const& max_dense_pts_size,
                          std::uint16_t const& dense_pts_length,
                          float* dense_pts_ptr){
  if(nullptr == dense_pts_ptr){
    std::cout << "NULL pointer" << std::endl;
    return filled_size;
  }
  if (ordered_pts_filled_size < 3U){
    return filled_size;
  }

  float cur_x, cur_y, cw1, cw2;
  std::uint16_t filled_cur_line{0U};
  for (std::uint16_t i{1U}; i < ordered_pts_filled_size - 1U; ++i){
    if (filled_size > max_dense_pts_size){
      break;
    }

    const Point& fst_pts = ordered_pts.at(i);
    const Point& sec_pts = ordered_pts.at(i+1);

    cur_x = std::min(std::floor(fst_pts.x+EPS), std::floor(sec_pts.x+EPS));
    cur_y = std::min(std::floor(fst_pts.y+EPS), std::floor(sec_pts.y+EPS));

    // 第一二个点在同一条边上，调整分配给相邻栅格
    if ( filled_cur_line == 0 && (std::fabs(fst_pts.x - sec_pts.x) < EPS ||  std::fabs(fst_pts.y - sec_pts.y) < EPS) &&
          ordered_pts_filled_size > 3U){
      const Point& next_pts = ordered_pts.at(i+2);
      if(std::fabs(fst_pts.x - sec_pts.x) < EPS && next_pts.x > sec_pts.x + EPS &&
         std::fabs(fst_pts.x - std::floor(fst_pts.x+EPS)) < EPS){
        cur_x = cur_x - 1.0;
      } else if (std::fabs(fst_pts.y - sec_pts.y) < EPS && next_pts.y > sec_pts.y + EPS  &&
                 std::fabs(fst_pts.y - std::floor(fst_pts.y+EPS)) < EPS){
          cur_y = cur_y - 1.0;
        }
    }

    if(filled_size > 0 && filled_cur_line > 0){
      float const& last_x = dense_pts_ptr[(filled_size-1) * dense_pts_length + 2];
      float const& last_y = dense_pts_ptr[(filled_size-1) * dense_pts_length + 3];

      // 中断分为两条线
      if((std::fabs(cur_x -  last_x) + std::fabs(cur_y - last_y)) > 3.1f){
        id++;
      }

      // 两点在同一条边上，调整分配给相邻栅格
      if(std::fabs(std::fabs(cur_x - last_x) + std::fabs(cur_y - last_y) - 1.0f) > EPS){
        if (std::fabs(fst_pts.x - sec_pts.x) < EPS){
          cur_x = cur_x - 1.0;
        } else if (std::fabs(fst_pts.y - sec_pts.y) < EPS){
          cur_y = cur_y - 1.0;
        }
      }
    }

    // x < 0 || x >= max_x || y < 0 || y >= max_y
    if (cur_x < (0 - EPS) || cur_x > (max_xyz_[0] - 1 + EPS) ||
        cur_y < (0 - EPS) || cur_y > (max_xyz_[1] - 1 + EPS)){
      continue;
    }

    dense_pts_ptr[filled_size * dense_pts_length + 0] = id;
    dense_pts_ptr[filled_size * dense_pts_length + 1] = 0.0;
    dense_pts_ptr[filled_size * dense_pts_length + 2] = cur_x;
    dense_pts_ptr[filled_size * dense_pts_length + 3] = cur_y;
    dense_pts_ptr[filled_size * dense_pts_length + 4] = type;
    dense_pts_ptr[filled_size * dense_pts_length + 7] = cls;

    cw1 = get_clockwise(fst_pts.x - cur_x, fst_pts.y - cur_y);
    cw2 = get_clockwise(sec_pts.x - cur_x, sec_pts.y - cur_y);

    dense_pts_ptr[filled_size * dense_pts_length + 5] = std::min(cw1, cw2);
    dense_pts_ptr[filled_size * dense_pts_length + 6] = 4.0 - std::max(cw1, cw2);
    filled_size ++;
    filled_cur_line++;
  }
  id++;
  return filled_size;
}

std::uint16_t process_one_line(float const* gt_lines_pts_ptr,
                               std::uint16_t const& start_i,
                               std::uint16_t const& end_i,
                               std::uint16_t const& gt_lines_length,
                               std::array<float, 3> const& voxels_size_,
                               std::array<float, 6> const& pts_range_,
                               std::array<Point, 2048>& ordered_pts) {
  if(nullptr == gt_lines_pts_ptr){
    std::cout << "NULL pointer" << std::endl;
  }

//  std::cout << "=============================================" << std::endl;
//  std::cout << "start_i, end_i: " << start_i << ", " << end_i << std::endl;


  std::array<Point, 2048> x_pts, y_pts;
  std::uint16_t x_filled_size{0U}, y_filled_size{0U}, ordered_pts_filled_size{1U};

  Point intersect_p, p0, p1, q0, q1;
  bool flag, crossed;
  float start_x, start_y, end_x, end_y;

  for (std::uint16_t i{start_i}; i < end_i; ++i){
    x_filled_size = 0U;
    y_filled_size = 0U;

    start_x = (gt_lines_pts_ptr[i * gt_lines_length + 1] - pts_range_[0]) / voxels_size_[0];
    start_y = (gt_lines_pts_ptr[i * gt_lines_length + 2] - pts_range_[1]) / voxels_size_[1];
    end_x = (gt_lines_pts_ptr[(i + 1) * gt_lines_length + 1] - pts_range_[0]) / voxels_size_[0];
    end_y = (gt_lines_pts_ptr[(i + 1) * gt_lines_length + 2] - pts_range_[1]) / voxels_size_[1];

//    std::cout << "===================== " << i << " ========================" << std::endl;
//    std::cout << "start_x, start_y: " << start_x << ", " << start_y << std::endl;
//    std::cout << "end_x, end_y: " << end_x << ", " << end_y << std::endl;

    p0.x = start_x;
    p0.y = start_y;
    p1.x = end_x;
    p1.y = end_y;

    ordered_pts.at(0).x = p0.x;
    ordered_pts.at(0).y = p0.y;

    // calculate cross pts with y-axis
    flag = false;
    if (start_x > end_x){
      std::swap(start_x, end_x);
      std::swap(start_y, end_y);
      flag = !flag;         //flag: true
    }

    for(float cur_x{std::ceil(start_x)}; cur_x < end_x; ++cur_x){
      q0.x = cur_x;
      q0.y = pts_range_[1];
      q1.x = cur_x;
      q1.y = pts_range_[4];
      crossed = get_cross_pts(p0, p1, q0, q1, intersect_p);
//      std::cout << "cur_x, pts_range_[1]: " << cur_x << ", " << pts_range_[1] << std::endl;
//      std::cout << "cur_x, pts_range_[4]: " << cur_x << ", " << pts_range_[4] << std::endl;

      if (crossed && (x_filled_size < x_pts.size())){
        intersect_p.x = std::roundf(intersect_p.x);

        x_pts.at(x_filled_size).x = intersect_p.x;
        x_pts.at(x_filled_size).y = intersect_p.y;
        x_filled_size++;
//        std::cout << "intersect_p.x, intersect_p.y: " << intersect_p.x << ", " << intersect_p.y << std::endl;
      }
    }

    if(flag){
      std::reverse(x_pts.begin(), x_pts.begin()+x_filled_size);
    }


    // calculate cross pts with x-axis
    if (start_y > end_y){
      std::swap(start_x, end_x);
      std::swap(start_y, end_y);
      flag = !flag;
    }

    for(float cur_y{std::ceil(start_y)}; cur_y < end_y; ++cur_y){
      q0.x = pts_range_[0];
      q0.y = cur_y;
      q1.x = pts_range_[3];
      q1.y = cur_y;

      crossed = get_cross_pts(p0, p1, q0, q1, intersect_p);
//      std::cout << "pts_range_[0], cur_y: " << pts_range_[0] << ", " << cur_y << std::endl;
//      std::cout << "pts_range_[3], cur_y: " << pts_range_[3] << ", " << cur_y << std::endl;
      if (crossed && (y_filled_size < y_pts.size())){
        intersect_p.y = std::roundf(intersect_p.y);

        y_pts.at(y_filled_size).x = intersect_p.x;
        y_pts.at(y_filled_size).y = intersect_p.y;
        y_filled_size++;
//        std::cout << "intersect_p.x, intersect_p.y: " << intersect_p.x << ", " << intersect_p.y << std::endl;
      }
    }

    if(flag){
      std::reverse(y_pts.begin(), y_pts.begin()+y_filled_size);
    }

    // sort cross pts
    ordered_pts_filled_size = sort_cross_pts(ordered_pts, x_pts, y_pts,
                                              ordered_pts_filled_size, x_filled_size, y_filled_size);
  }
  return ordered_pts_filled_size;
}

std::uint16_t sort_cross_pts(std::array<Point, 2048>& ordered_pts,
                             std::array<Point, 2048> const& x_pts,
                             std::array<Point, 2048> const& y_pts,
                             std::uint16_t ordered_pts_filled_size,
                             std::uint16_t const& x_filled_size,
                             std::uint16_t const& y_filled_size){
  if(ordered_pts_filled_size == 0U){
    std::cout << "Wrong searching start!" << std::endl;
    return ordered_pts_filled_size;
  }

//  std::cout << "===================== x_pts ========================" << std::endl;
//  for (std::uint16_t i{0U}; i < x_filled_size; ++i)
//  {
//      const Point& cur_pts = x_pts.at(i);
//      std::cout << "i: " << i << std::endl;
//      std::cout << "cur_pts.x: " << cur_pts.x << std::endl;
//      std::cout << "cur_pts.y: " << cur_pts.y << std::endl;
//  }
//  std::cout << "===================== y_pts ========================" << std::endl;
//  for (std::uint16_t i{0U}; i < y_filled_size; ++i)
//  {
//      const Point& cur_pts = y_pts.at(i);
//      std::cout << "i: " << i << std::endl;
//      std::cout << "cur_pts.x: " << cur_pts.x << std::endl;
//      std::cout << "cur_pts.y: " << cur_pts.y << std::endl;
//  }
//
//  std::cout << "=============================================" << std::endl;

  float dis_x{0.0f}, dis_y{0.0f};
  bool first{true};
  std::uint16_t x_pos{0U}, y_pos{0U};

  while(x_pos < x_filled_size || y_pos < y_filled_size){
    if(x_pos == x_filled_size){
//      std::cout << "===================== insert y_pts ========================" << std::endl;
//      std::uint16_t count{0};
//      for (std::uint16_t i{y_pos}; i < y_filled_size; ++i)
//      {
//          const Point& cur_pts = y_pts.at(i);
//          std::cout << "===================== ordered_pts index: " << ordered_pts_filled_size + count << " ========================" << std::endl;
//          std::cout << "i: " << count << std::endl;
//          std::cout << "cur_pts.x: " << cur_pts.x << std::endl;
//          std::cout << "cur_pts.y: " << cur_pts.y << std::endl;
//          count++;
//      }

      for(; y_pos < y_filled_size; y_pos++){
        if (ordered_pts_filled_size < ordered_pts.size()){
          ordered_pts.at(ordered_pts_filled_size).x = y_pts.at(y_pos).x;
          ordered_pts.at(ordered_pts_filled_size).y = y_pts.at(y_pos).y;
          ordered_pts_filled_size++;
        }
      }
      break;
    }

    if(y_pos == y_filled_size){
//      std::cout << "===================== insert x_pts ========================" << std::endl;
//      std::uint16_t count{0};
//      for (std::uint16_t i{x_pos}; i < x_filled_size; ++i)
//      {
//          const Point& cur_pts = x_pts.at(i);
//          std::cout << "===================== ordered_pts index: " << ordered_pts_filled_size + count << " ========================" << std::endl;
//          std::cout << "i: " << count << std::endl;
//          std::cout << "cur_pts.x: " << cur_pts.x << std::endl;
//          std::cout << "cur_pts.y: " << cur_pts.y << std::endl;
//          count++;
//      }

      for(; x_pos < x_filled_size; x_pos++){
        if (ordered_pts_filled_size < ordered_pts.size()){
          ordered_pts.at(ordered_pts_filled_size).x = x_pts.at(x_pos).x;
          ordered_pts.at(ordered_pts_filled_size).y = x_pts.at(x_pos).y;
          ordered_pts_filled_size++;
        }
      }
      break;
    }

    if (!(ordered_pts_filled_size < ordered_pts.size())){
      break;
    }

    Point& last_pt =  first ? ordered_pts.at(0) : ordered_pts.at(ordered_pts_filled_size-1);
    Point const& x_pt =  x_pts.at(x_pos);
    Point const& y_pt =  y_pts.at(y_pos);
    dis_x = get_distance_square(last_pt, x_pt);
    dis_y = get_distance_square(last_pt, y_pt);

    first = false;
//    std::cout << "===================== ordered_pts index: " << ordered_pts_filled_size << " ========================" << std::endl;
//    std::cout << "last_pt.x, last_pt.y: " << last_pt.x << ", " << last_pt.y << std::endl;
//    std::cout << "x_pt.x, x_pt.y: " << x_pt.x << ", " << x_pt.y << std::endl;
//    std::cout << "y_pt.x, y_pt.y: " << y_pt.x << ", " << y_pt.y << std::endl;
//    std::cout << "dis_x: " << dis_x << std::endl;
//    std::cout << "dis_y: " << dis_y << std::endl;

    if (dis_x > dis_y){
      ordered_pts.at(ordered_pts_filled_size).x = y_pt.x;
      ordered_pts.at(ordered_pts_filled_size).y = y_pt.y;
      y_pos++;
    }else{
      ordered_pts.at(ordered_pts_filled_size).x = x_pt.x;
      ordered_pts.at(ordered_pts_filled_size).y = x_pt.y;
      x_pos++;
    }
    ordered_pts_filled_size++;
  }
  return ordered_pts_filled_size;
}

int gt_lines_voxel_interpolation(py::array_t<float> gt_lines_pts,
                                 py::array_t<float> voxels_size,
                                 py::array_t<float> pts_range,
                                 int max_dense_pts_size,
                                 py::array_t<float> dense_pts){
  // Get  ptr from python
  py::buffer_info gt_lines_pts_buf = gt_lines_pts.request();
  py::buffer_info voxels_size_buf = voxels_size.request();
  py::buffer_info pts_range_buf = pts_range.request();
  py::buffer_info dense_pts_buf = dense_pts.request();

  std::uint16_t gt_lines_num = gt_lines_pts_buf.shape[0];
  std::uint16_t gt_lines_length = gt_lines_pts_buf.shape[1];
  std::uint16_t dense_pts_length = dense_pts_buf.shape[1];
  std::uint16_t voxels_size_dim = voxels_size_buf.shape[0];


  float* gt_lines_pts_ptr = (float*)gt_lines_pts_buf.ptr;
  float* voxels_size_ptr = (float*)voxels_size_buf.ptr;
  float* pts_range_ptr = (float*)pts_range_buf.ptr;
  float* dense_pts_ptr = (float*)dense_pts_buf.ptr;

  std::uint16_t max_dense_pts_size_ = static_cast<std::uint16_t>(max_dense_pts_size);
  int filled_size = voxel_encode(gt_lines_pts_ptr, gt_lines_num, gt_lines_length, voxels_size_ptr, pts_range_ptr,
                                  max_dense_pts_size_, dense_pts_length, voxels_size_dim, dense_pts_ptr);
  return filled_size;
}
